package com.algorithm.greedy;

import java.util.Arrays;

/**
 * 普利姆算法
 * @author create by xcb 2025/3/6 下午8:43
 */
public class Prim {

    // 图的顶点数
    private static final int V = 6;

    // 查找最小权重的边
    private int minKey(int[] key, boolean[] mstSet) {
        int min = Integer.MAX_VALUE, minIndex = -1;

        for (int v = 0; v < V; v++) {
            if (!mstSet[v] && key[v] < min) {
                min = key[v];
                minIndex = v;
            }
        }
        return minIndex;
    }

    // 打印生成的最小生成树
    private void printMST(int[] parent, int[][] graph) {
        System.out.println("边 \t 权重");
        for (int i = 1; i < V; i++) {
            System.out.println(parent[i] + " - " + i + "\t " + graph[i][parent[i]]);
        }
    }

    // 实现 Prim 算法
    public void primMST(int[][] graph) {
        // 存储生成树
        int[] parent = new int[V];
        // 存储权重值
        int[] key = new int[V];
        // 标记已包含在 MST 中的顶点
        boolean[] mstSet = new boolean[V];

        // 所有键初始化为无穷大
        Arrays.fill(key, Integer.MAX_VALUE);
        // 第一个顶点的权重设置为 0
        key[0] = 0;
        // 第一条边不需要
        parent[0] = -1;

        for (int count = 0; count < V - 1; count++) {
            // 选择权重最小的顶点
            int u = minKey(key, mstSet);
            // 将选择的顶点标记为已包含在 MST 中
            mstSet[u] = true;

            for (int v = 0; v < V; v++) {
                // 更新与 u 连接的顶点 v 的键值和父节点
                if (graph[u][v] != 0 && !mstSet[v] && graph[u][v] < key[v]) {
                    // u 是 v 的 parent
                    parent[v] = u;
                    // 更新权重
                    key[v] = graph[u][v];
                }
            }
        }

        // 打印生成的最小生成树
        printMST(parent, graph);
    }

    public static void main(String[] args) {
        Prim mst = new Prim();
        int[][] graph = new int[][] {
                {0, 7, 9, 0, 0, 14},
                {7, 0, 10, 15, 0, 0},
                {9, 10, 0, 11, 0, 2},
                {0, 15, 11, 0, 6, 0},
                {0, 0, 0, 6, 0, 9},
                {14, 0, 2, 0, 9, 0}
        };
        mst.primMST(graph);
    }

}
